{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data='./input/data/train'\n",
    "validation_data='./input/data/validation'\n",
    "test_data='./input/data/test'\n",
    "local_model_dir='./model/'\n",
    "num_gpus=0\n",
    "input_data_config='File'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ['SM_CHANNEL_TRAIN']=train_data\n",
    "# os.environ['SM_CHANNEL_VALIDATION']=validation_data\n",
    "# os.environ['SM_CHANNEL_TEST']=test_data\n",
    "# os.environ['SM_MODEL_DIR']=local_model_dir\n",
    "# os.environ['SM_NUM_GPUS']=str(num_gpus)\n",
    "# os.environ['SM_INPUT_DATA_CONFIG']=input_data_config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-Tuning a BERT Model and Create a Text Classifier\n",
    "\n",
    "We have already performed the Feature Engineering to create BERT embeddings from the `reviews_body` text using the pre-trained BERT model, and split the dataset into train, validation and test files. To optimize for Tensorflow training, we saved the files in TFRecord format. \n",
    "\n",
    "Now, let’s fine-tune the BERT model to our Customer Reviews Dataset and add a new classification layer to predict the `star_rating` for a given `review_body`.\n",
    "\n",
    "![BERT Training](img/bert_training.png)\n",
    "\n",
    "As mentioned earlier, BERT’s attention mechanism is called a Transformer. This is, not coincidentally, the name of the popular BERT Python library, “Transformers,” maintained by a company called [HuggingFace](https://github.com/huggingface/transformers). We will use a variant of BERT called [DistilBert](https://arxiv.org/pdf/1910.01108.pdf) which requires less memory and compute, but maintains very good accuracy on our dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DEMO 1: \n",
    "# Develop Model Training Code In Noteboook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q easydict==1.9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q transformers==2.8.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q tensorflow==2.1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q scikit-learn==0.23.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs=1\n",
    "learning_rate=0.00001\n",
    "epsilon=0.00000001\n",
    "train_batch_size=8\n",
    "validation_batch_size=8\n",
    "test_batch_size=8\n",
    "train_steps_per_epoch=1\n",
    "validation_steps=1\n",
    "test_steps=1\n",
    "use_xla=True\n",
    "use_amp=False\n",
    "max_seq_length=64\n",
    "freeze_bert_layer=True\n",
    "run_validation=True\n",
    "run_test=True\n",
    "run_sample_predictions=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import random\n",
    "import pandas as pd\n",
    "import easydict\n",
    "from glob import glob\n",
    "import pprint\n",
    "import argparse\n",
    "import json\n",
    "import subprocess\n",
    "import sys\n",
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "from transformers import DistilBertTokenizer\n",
    "from transformers import TFDistilBertForSequenceClassification\n",
    "from transformers import TextClassificationPipeline\n",
    "from transformers.configuration_distilbert import DistilBertConfig\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint\n",
    "from tensorflow.keras.models import load_model\n",
    "\n",
    "CLASSES = [1, 2, 3, 4, 5]\n",
    "\n",
    "def select_data_and_label_from_record(record):\n",
    "    x = {\n",
    "        'input_ids': record['input_ids'],\n",
    "        'input_mask': record['input_mask'],\n",
    "        'segment_ids': record['segment_ids']\n",
    "    }\n",
    "\n",
    "    y = record['label_ids']\n",
    "\n",
    "    return (x, y)\n",
    "\n",
    "\n",
    "def file_based_input_dataset_builder(channel,\n",
    "                                     input_filenames,\n",
    "                                     pipe_mode,\n",
    "                                     is_training,\n",
    "                                     drop_remainder,\n",
    "                                     batch_size,\n",
    "                                     epochs,\n",
    "                                     steps_per_epoch,\n",
    "                                     max_seq_length):\n",
    "\n",
    "    # For training, we want a lot of parallel reading and shuffling.\n",
    "    # For eval, we want no shuffling and parallel reading doesn't matter.\n",
    "\n",
    "    if pipe_mode:\n",
    "            print('***** Using pipe_mode with channel {}'.format(channel))\n",
    "            from sagemaker_tensorflow import PipeModeDataset\n",
    "            dataset = PipeModeDataset(channel=channel,\n",
    "                                      record_format='TFRecord')\n",
    "    else:\n",
    "        print('***** Using input_filenames {}'.format(input_filenames))\n",
    "        dataset = tf.data.TFRecordDataset(input_filenames)\n",
    "    \n",
    "    dataset = dataset.repeat(epochs * steps_per_epoch * 100)\n",
    "\n",
    "    name_to_features = {\n",
    "      \"input_ids\": tf.io.FixedLenFeature([max_seq_length], tf.int64),\n",
    "      \"input_mask\": tf.io.FixedLenFeature([max_seq_length], tf.int64),\n",
    "      \"segment_ids\": tf.io.FixedLenFeature([max_seq_length], tf.int64),\n",
    "      \"label_ids\": tf.io.FixedLenFeature([], tf.int64),\n",
    "    }\n",
    "\n",
    "    def _decode_record(record, name_to_features):\n",
    "        \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n",
    "        record = tf.io.parse_single_example(record, name_to_features)\n",
    "        return record\n",
    "    \n",
    "    dataset = dataset.apply(\n",
    "        tf.data.experimental.map_and_batch(\n",
    "          lambda record: _decode_record(record, name_to_features),\n",
    "          batch_size=batch_size,\n",
    "          drop_remainder=drop_remainder,\n",
    "          num_parallel_calls=tf.data.experimental.AUTOTUNE))\n",
    "\n",
    "    dataset = dataset.shuffle(buffer_size=1000,\n",
    "                              reshuffle_each_iteration=True)\n",
    "    \n",
    "    row_count = 0\n",
    "    print('**************** {} *****************'.format(channel))\n",
    "    for row in dataset.as_numpy_iterator():\n",
    "        if row_count == 1:\n",
    "            break\n",
    "        print(row)\n",
    "        row_count = row_count + 1\n",
    "    return dataset\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "\n",
    "    args=easydict.EasyDict({\n",
    "        'train_data': train_data,\n",
    "        'validation_data': validation_data,\n",
    "        'test_data': test_data,\n",
    "        'local_model_dir': local_model_dir,\n",
    "        'num_gpus': num_gpus,\n",
    "        'use_xla': use_xla,\n",
    "        'use_amp': use_amp,\n",
    "        'max_seq_length': max_seq_length,\n",
    "        'train_batch_size': train_batch_size,\n",
    "        'validation_batch_size': validation_batch_size,\n",
    "        'test_batch_size': test_batch_size,\n",
    "        'epochs': epochs,\n",
    "        'learning_rate': learning_rate,\n",
    "        'epsilon': epsilon,\n",
    "        'train_steps_per_epoch': train_steps_per_epoch,\n",
    "        'validation_steps': validation_steps,\n",
    "        'test_steps': test_steps,\n",
    "        'freeze_bert_layer': freeze_bert_layer,\n",
    "        'run_validation': run_validation,\n",
    "        'run_test': run_test,\n",
    "        'run_sample_predictions': run_sample_predictions,\n",
    "        'input_data_config': input_data_config\n",
    "    })\n",
    "    \n",
    "    \n",
    "    env_var = os.environ \n",
    "    print(\"Environment Variables:\") \n",
    "    pprint.pprint(dict(env_var), width = 1) \n",
    "    \n",
    "    train_data = args.train_data\n",
    "    print('train_data {}'.format(train_data))\n",
    "    validation_data = args.validation_data\n",
    "    print('validation_data {}'.format(validation_data))\n",
    "    test_data = args.test_data\n",
    "    print('test_data {}'.format(test_data))    \n",
    "    local_model_dir = args.model_dir\n",
    "    print('local_model_dir {}'.format(local_model_dir))         \n",
    "    num_gpus = args.num_gpus\n",
    "    print('num_gpus {}'.format(num_gpus))   \n",
    "    use_xla = args.use_xla\n",
    "    print('use_xla {}'.format(use_xla))    \n",
    "    use_amp = args.use_amp\n",
    "    print('use_amp {}'.format(use_amp))    \n",
    "    max_seq_length = args.max_seq_length\n",
    "    print('max_seq_length {}'.format(max_seq_length))    \n",
    "    train_batch_size = args.train_batch_size\n",
    "    print('train_batch_size {}'.format(train_batch_size))    \n",
    "    validation_batch_size = args.validation_batch_size\n",
    "    print('validation_batch_size {}'.format(validation_batch_size))    \n",
    "    test_batch_size = args.test_batch_size\n",
    "    print('test_batch_size {}'.format(test_batch_size))    \n",
    "    epochs = args.epochs\n",
    "    print('epochs {}'.format(epochs))    \n",
    "    learning_rate = args.learning_rate\n",
    "    print('learning_rate {}'.format(learning_rate))    \n",
    "    epsilon = args.epsilon\n",
    "    print('epsilon {}'.format(epsilon))    \n",
    "    train_steps_per_epoch = args.train_steps_per_epoch\n",
    "    print('train_steps_per_epoch {}'.format(train_steps_per_epoch))    \n",
    "    validation_steps = args.validation_steps\n",
    "    print('validation_steps {}'.format(validation_steps))    \n",
    "    test_steps = args.test_steps\n",
    "    print('test_steps {}'.format(test_steps))    \n",
    "    freeze_bert_layer = args.freeze_bert_layer\n",
    "    print('freeze_bert_layer {}'.format(freeze_bert_layer))    \n",
    "    run_validation = args.run_validation\n",
    "    print('run_validation {}'.format(run_validation))    \n",
    "    run_test = args.run_test\n",
    "    print('run_test {}'.format(run_test))    \n",
    "    run_sample_predictions = args.run_sample_predictions\n",
    "    print('run_sample_predictions {}'.format(run_sample_predictions))\n",
    "    input_data_config = args.input_data_config\n",
    "    print('input_data_config {}'.format(input_data_config))\n",
    "    \n",
    "    # Determine if PipeMode is enabled \n",
    "    pipe_mode = (input_data_config.find('Pipe') >= 0)\n",
    "    print('Using pipe_mode: {}'.format(pipe_mode))\n",
    " \n",
    "    # Model Output \n",
    "    transformer_fine_tuned_model_path = os.path.join(local_model_dir, 'transformers/fine-tuned/')\n",
    "    os.makedirs(transformer_fine_tuned_model_path, exist_ok=True)\n",
    "\n",
    "    # SavedModel Output\n",
    "    tensorflow_saved_model_path = os.path.join(local_model_dir, 'tensorflow/saved_model/0')\n",
    "    os.makedirs(tensorflow_saved_model_path, exist_ok=True) \n",
    "    \n",
    "    distributed_strategy = tf.distribute.MirroredStrategy()\n",
    "\n",
    "    with distributed_strategy.scope():\n",
    "        tf.config.optimizer.set_jit(use_xla)\n",
    "        tf.config.optimizer.set_experimental_options({\"auto_mixed_precision\": use_amp})\n",
    "\n",
    "        train_data_filenames = glob(os.path.join(train_data, '*.tfrecord'))\n",
    "        print('train_data_filenames {}'.format(train_data_filenames))\n",
    "        train_dataset = file_based_input_dataset_builder(\n",
    "            channel='train',\n",
    "            input_filenames=train_data_filenames,\n",
    "            pipe_mode=pipe_mode,\n",
    "            is_training=True,\n",
    "            drop_remainder=False,\n",
    "            batch_size=train_batch_size,\n",
    "            epochs=epochs,\n",
    "            steps_per_epoch=train_steps_per_epoch,\n",
    "            max_seq_length=max_seq_length).map(select_data_and_label_from_record)\n",
    "\n",
    "        tokenizer = None\n",
    "        config = None\n",
    "        model = None\n",
    "\n",
    "        successful_download = False\n",
    "        retries = 0\n",
    "        while (retries < 5 and not successful_download):\n",
    "            try:\n",
    "                tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "                config = DistilBertConfig.from_pretrained('distilbert-base-uncased',\n",
    "                                                          num_labels=len(CLASSES))\n",
    "                model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased',\n",
    "                                                                              config=config)\n",
    "                successful_download = True\n",
    "                print('Sucessfully downloaded after {} retries.'.format(retries))\n",
    "            except:\n",
    "                retries = retries + 1\n",
    "                random_sleep = random.randint(1, 30)\n",
    "                print('Retry #{}.  Sleeping for {} seconds'.format(retries, random_sleep))\n",
    "                time.sleep(random_sleep)\n",
    "\n",
    "        callbacks = []       \n",
    "        initial_epoch_number = 0 \n",
    "\n",
    "        if not tokenizer or not model or not config:\n",
    "            print('Not properly initialized...')\n",
    "\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate, epsilon=epsilon)\n",
    "        print('** use_amp {}'.format(use_amp))        \n",
    "        if use_amp:\n",
    "            # loss scaling is currently required when using mixed precision\n",
    "            optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(optimizer, 'dynamic')\n",
    "  \n",
    "        print('*** OPTIMIZER {} ***'.format(optimizer))\n",
    "        \n",
    "        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "        metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')\n",
    "\n",
    "        model.compile(optimizer=optimizer, loss=loss, metrics=[metric])\n",
    "        print('Compiled model {}'.format(model))          \n",
    "        model.layers[0].trainable = not freeze_bert_layer\n",
    "        print(model.summary())\n",
    "\n",
    "        if run_validation:\n",
    "            validation_data_filenames = glob(os.path.join(validation_data, '*.tfrecord'))\n",
    "            print('validation_data_filenames {}'.format(validation_data_filenames))\n",
    "            validation_dataset = file_based_input_dataset_builder(\n",
    "                channel='validation',\n",
    "                input_filenames=validation_data_filenames,\n",
    "                pipe_mode=pipe_mode,\n",
    "                is_training=False,\n",
    "                drop_remainder=False,\n",
    "                batch_size=validation_batch_size,\n",
    "                epochs=epochs,\n",
    "                steps_per_epoch=validation_steps,\n",
    "                max_seq_length=max_seq_length).map(select_data_and_label_from_record)\n",
    "            \n",
    "            print('Starting Training and Validation...')\n",
    "            validation_dataset = validation_dataset.take(validation_steps)\n",
    "            train_and_validation_history = model.fit(train_dataset,\n",
    "                                                     shuffle=True,\n",
    "                                                     epochs=epochs,\n",
    "                                                     initial_epoch=initial_epoch_number,\n",
    "                                                     steps_per_epoch=train_steps_per_epoch,\n",
    "                                                     validation_data=validation_dataset,\n",
    "                                                     validation_steps=validation_steps,\n",
    "                                                     callbacks=callbacks)                                \n",
    "            print(train_and_validation_history)\n",
    "        else: # Not running validation\n",
    "            print('Starting Training (Without Validation)...')\n",
    "            train_history = model.fit(train_dataset,\n",
    "                                      shuffle=True,\n",
    "                                      epochs=epochs,\n",
    "                                      initial_epoch=initial_epoch_number,\n",
    "                                      steps_per_epoch=train_steps_per_epoch,\n",
    "                                      callbacks=callbacks)                \n",
    "            print(train_history)\n",
    "\n",
    "        if run_test:\n",
    "            test_data_filenames = glob(os.path.join(test_data, '*.tfrecord'))\n",
    "            print('test_data_filenames {}'.format(test_data_filenames))\n",
    "            test_dataset = file_based_input_dataset_builder(\n",
    "                channel='test',\n",
    "                input_filenames=test_data_filenames,\n",
    "                pipe_mode=pipe_mode,\n",
    "                is_training=False,\n",
    "                drop_remainder=False,\n",
    "                batch_size=test_batch_size,\n",
    "                epochs=epochs,\n",
    "                steps_per_epoch=test_steps,\n",
    "                max_seq_length=max_seq_length).map(select_data_and_label_from_record)\n",
    "\n",
    "            print('Starting test...')\n",
    "            test_history = model.evaluate(test_dataset,\n",
    "                                          steps=test_steps,\n",
    "                                          callbacks=callbacks)\n",
    "                                 \n",
    "            print('Test history {}'.format(test_history))\n",
    "            \n",
    "        # Save the Fine-Tuned Transformers Model as a New \"Pre-Trained\" Model\n",
    "        print('transformer_fine_tuned_model_path {}'.format(transformer_fine_tuned_model_path))   \n",
    "        model.save_pretrained(transformer_fine_tuned_model_path)\n",
    "\n",
    "        # Save the TensorFlow SavedModel for Serving Predictions\n",
    "        print('tensorflow_saved_model_path {}'.format(tensorflow_saved_model_path))   \n",
    "        model.save(tensorflow_saved_model_path, save_format='tf')\n",
    "                \n",
    "    if run_sample_predictions:\n",
    "        loaded_model = TFDistilBertForSequenceClassification.from_pretrained(transformer_fine_tuned_model_path,\n",
    "                                                                       id2label={\n",
    "                                                                        0: 1,\n",
    "                                                                        1: 2,\n",
    "                                                                        2: 3,\n",
    "                                                                        3: 4,\n",
    "                                                                        4: 5\n",
    "                                                                       },\n",
    "                                                                       label2id={\n",
    "                                                                        1: 0,\n",
    "                                                                        2: 1,\n",
    "                                                                        3: 2,\n",
    "                                                                        4: 3,\n",
    "                                                                        5: 4\n",
    "                                                                       })\n",
    "\n",
    "        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "\n",
    "        if num_gpus >= 1:\n",
    "            inference_device = 0 # GPU 0\n",
    "        else:\n",
    "            inference_device = -1 # CPU\n",
    "        print('inference_device {}'.format(inference_device))\n",
    "\n",
    "        inference_pipeline = TextClassificationPipeline(model=loaded_model, \n",
    "                                                        tokenizer=tokenizer,\n",
    "                                                        framework='tf',\n",
    "                                                        device=inference_device)  \n",
    "\n",
    "        print(\"\"\"I loved it!  I will recommend this to everyone.\"\"\", inference_pipeline(\"\"\"I loved it!  I will recommend this to everyone.\"\"\"))\n",
    "        print(\"\"\"It's OK.\"\"\", inference_pipeline(\"\"\"It's OK.\"\"\"))\n",
    "        print(\"\"\"Really bad.  I hope they don't make this anymore.\"\"\", inference_pipeline(\"\"\"Really bad.  I hope they don't make this anymore.\"\"\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
